from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging
import os
import sys

# configure logging
logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                    level=logging.INFO,
                    stream=sys.stdout)

import numpy as np
import scipy as scp
import scipy.misc
import tensorflow as tf

flags = tf.app.flags
FLAGS = flags.FLAGS
sys.path.insert(0, 'homo_net')
# sys.path.insert(1, 'incl')
from seg_utils import seg_utils as seg

try:
    # check whether setup was done correctly
    import tensorvision.utils as tv_utils
    import tensorvision.core as core
except ImportError:
    # You forgot to initialize submodules
    logging.error("Could not import the submodules. ")
    # logging.error("Please execute: "
    #               "'git submodule update --init --recursive")
    exit(1)

flags.DEFINE_string('logdir', None, 'Path to logdir.')
flags.DEFINE_string('input_image', None, 'Image to apply KittiSeg.')
flags.DEFINE_string('output_image', None, 'Image to apply KittiSeg.')

default_run = 'KittiSeg_pretrained'
weights_url = ("ftp://mi.eng.cam.ac.uk/pub/mttt2/models/KittiSeg_pretrained.zip")

def maybe_download_and_extract(runs_dir):
    logdir = os.path.join(runs_dir, default_run)
    if os.path.exists(logdir):
        # weights are download. Nothing to do
        return

    import zipfile
    download_name = tv_utils.download(weights_url, runs_dir)
    logging.info("Extracting KittiSeg_pretrained.zip")
    zipfile.ZipFile(download_name, 'r').extractall(runs_dir)
    return

def resize_label_image(image, gt_image, image_height, image_width):
    image = scp.misc.imresize(image, size=(image_height, image_width), interp='cubic')
    shape = gt_image.shape
    gt_image = scp.misc.imresize(gt_image, size=(image_height, image_width), interp='nearest')
    return image, gt_image

def main(_):
    tv_utils.set_gpus_to_use()
    if FLAGS.input_image is None:
        logging.error("No input_images was given.")
        logging.info(
            "Usage: python demo.py --input_image data/test.png "
            "[--output_image output_image] [--logdir /path/to/weights] "
            "[--gpus GPUs_to_use] ")
        exit(1)
    if FLAGS.logdir is None:
        # Download and use weights from the MultiNet Paper
        if 'TV_DIR_RUNS' in os.environ:
            runs_dir = os.path.join(os.environ['TV_DIR_RUNS'], 'KittiSeg')
        else:
            runs_dir = 'RUNS'
        maybe_download_and_extract(runs_dir)
        logdir = os.path.join(runs_dir, default_run)
    else:
        logging.info("Using weights found in{}".format(FLAGS.logdir))
        logdir = FLAGS.logdir

    # Loading hyperparameters from logdir
    hypes = tv_utils.load_hypes_from_logdir(logdir)
    logging.info("Hypes loaded successfully.")

    # Loading tv modules (encoder.py, decoder.py, eval.py) from logdir
    modules = tv_utils.load_modules_from_logdir(logdir)
    logging.info("Modules loaded successfully. Starting to build tf gragh.")

    # create tf gragh and build module
    with tf.Graph().as_default():
        # Create tf gragh and build module
        image_pl = tf.placeholder(tf.float32)
        image = tf.expand_dims(image_pl, 0)

        # build Tensorflow gragh and build module
        prediction = core.build_inference_graph(hypes, modules, image=image)
        logging.info("Gragh build successfully.")

        # create a session for running Ops on the gragh.
        sess = tf.Session()
        saver = tf.train.Saver()

        # Load weights from logdir
        core.load_weights(logdir, sess, saver)
        logging.info("Weights loaded successfully")

    input_image = FLAGS.input_image
    logging.info("Starting inference using{} as input".format(input_image))

    # Load and resize input image
    image = scp.misc.imread(input_image)
    if hypes['jitter']['resize_image']:
        # Resize innput only, if specified in hypes
        image_height = hypes['jitter']['image_height']
        image_width = hypes['jitter']['image_width']
        image = scp.misc.imresize(image, size=(image_height, image_width), interp='cubic')

    # Run KittiSeg model on image
    feed = {image_pl: image}
    softmax = prediction['softmax']
    output = sess.run([softmax], feed_dict=feed)

    # todo faster
    # for image in images:
    #     feed = {image_pl:image}
    #     softmax = prediction['softmax']
    #     output = sess.run([softmax], feed_dict=feed)

    # Reshape output from flat vector to 2D Image
    shape = image.shape
    output_image = output[0][:, 1].reshape(shape[0], shape[1])

    # Plot confidences as red-blue overlay
    rb_image = seg.make_overlay(image, output_image)

    # Accept all pixel with conf >= 0.5 as positive prediction
    # This creates a 'hard' prediction result for class street
    threshold = 0.5
    street_prediction = output_image > threshold
    # Plot the hard prediction as green overlay
    green_image = tv_utils.fast_overlay(image, street_prediction)

    # Save output imagaes to disk
    if FLAGS.output_image is None:
        output_base_name = input_image
    else:
        output_base_name = FLAGS.output_image
    raw_image_name = output_base_name.split('.')[0] + '_raw.png'
    rb_image_name = output_base_name.split('.')[0] + '_rb.png'
    green_image_name = output_base_name.split('.')[0] + '_green.png'
    scp.misc.imsave(raw_image_name, output_image)
    scp.misc.imsave(rb_image_name, rb_image)
    scp.misc.imsave(green_image_name, green_image)

    logging.info("Raw output image has been saved to: {}".format(os.path.realpath(raw_image_name)))
    logging.info("Red-Blue overlay of confs have been saved to: {}".format(os.path.realpath(rb_image_name)))
    logging.info("Green plot of predictions have been saved to: {}".format(os.path.realpath(green_image_name)))
    logging.warning("Do NOT use this Code to evaluate multiple images.")
    logging.warning("Demo.py is **very slow** and designed to be a tutorial to show how the KittiSeg works.")
    logging.warning("Please see this comment, if you like to apply demo.py to multiple images see:")
    logging.warning("https://github.com/MarvinTeichmann/KittiBox/issues/15#issuecomment-301800058")

if __name__ == '__main__':
    tf.app.run()